﻿module ScriptMonad

open Microsoft.Xna.Framework
open Microsoft.FSharp
open Microsoft.FSharp.Core
open System

type Script<'a,'s> = 's -> ScriptStep<'a,'s>
and ScriptStep<'a,'s> =
  Return of 'a
  | Yield of Script<'a,'s>
  | ArrowYield of Script<'a,'s>

type ScriptBuilder() =
  member this.Return(x:'a) : Script<'a,'s> =
    fun s -> Return x
  member this.Zero() : Script<Unit,'s> = this.Return()
  member this.Bind(p : Script<'a,'s>,
                   k : 'a -> Script<'b,'s>) : Script<'b,'s> =
    fun s ->
      match p s with
      | Return x -> k x s
      | Yield p' -> Yield(this.Bind(p',k))
      | ArrowYield p' -> ArrowYield(this.Bind(p',k))
  member this.Combine(p1:Script<'a,'s>,

                      p2:Script<'b,'s>) : Script<'b,'s> =
    this.Bind(p1, fun _ -> p2)
  member this.For(items:seq<'a>, body:'a -> Script<Unit,'s>) : Script<Unit,'s> =
    this.ForList(items |> Seq.toList, body)
  member this.ForList(items:list<'a>, body:'a -> Script<Unit,'s>) : Script<Unit,'s> =
    fun s ->
      match items with
      | [] -> Return()
      | hd :: tl -> 
        this.Combine(body hd, this.ForList(tl, body)) s
  member this.ReturnFrom(s:Script<'a,'s>) = s
  member this.Delay s = this.Bind(this.Return(), s)
  member this.Run s = s

let script = ScriptBuilder()

let arrow_yield_ : Script<Unit,'s> = fun s -> ArrowYield(fun s -> Return())
let yield_ : Script<Unit,'s> = fun s -> Yield(fun s -> Return())
let begin_ : Script<Unit,'s> = fun s -> Return()
let get_state = fun s -> Return s

let lift_1 (f:'a1 -> 'b) s1 =
  script{
    let! x1 = s1
    return f x1
  }

let lift_2 (f:'a1 -> 'a2 -> 'b) s1 s2 =
  script{
    let! x1 = s1
    let! x2 = s2
    return f x1 x2
  }

let not_ (s:Script<bool,'s>) : Script<bool,'s> =
  lift_1 not s

let and_ (s1:Script<bool,'s>) (s2:Script<bool,'s>) : Script<bool,'s> =
  lift_2 (&&) s1 s2

let or_ (s1:Script<bool,'s>) (s2:Script<bool,'s>) : Script<bool,'s> =
  lift_2 (||) s1 s2

let ignore_ (s:Script<'a,'s>) : Script<Unit,'s> =
  lift_1 (fun x -> ()) s

let rec (.&&) (s1:Script<'a,'s>) (s2:Script<'b,'s>) : Script<'a * 'b,'s> =
  fun s ->
    match s1 s,s2 s with
    | Return x, Return y    -> Return (x,y)
    | Yield k1, Yield k2    -> (.&&) k1 k2 |> Yield
    | Yield k1, Return y    -> (.&&) k1 (fun s -> Return y) |> Yield
    | Return x, Yield k2    -> (.&&) (fun s -> Return x) k2 |> Yield
    | Yield k1, ArrowYield k2 -> (.&&) k1 k2 |> Yield
    | ArrowYield k1, Yield k2 -> (.&&) k1 k2 |> Yield
    | ArrowYield k1, ArrowYield k2 -> (.&&) k1 k2 |> Yield
    | ArrowYield k1, Return y    -> (.&&) k1 (fun s -> Return y) |> Yield
    | Return x, ArrowYield k2    -> (.&&) (fun s -> Return x) k2 |> Yield

let (.&&>) s1 s2 = s1 .&& s2 |> ignore_
let parallel_ = (.&&)

let rec (.&&&) =
  function
  | [] -> script{ return() }
  | x::xs -> parallel_ x (xs |> (.&&&)) |> ignore_

let parallel_many_ = (.&&&)

let rec (.||) (s1:Script<'a,'s>) (s2:Script<'b,'s>) : Script<Choice<'a,'b>,'s> =
  fun s ->
    match s1 s,s2 s with
    | Return x, _        -> Return(Choice1Of2 x)
    | _, Return y        -> Return(Choice2Of2 y)
    | ArrowYield k1, _   ->
      script{
        let! res = k1
        return Choice1Of2 res
      } |> Yield
    | _, ArrowYield k2   ->
      script{
        let! res = k2
        return Choice2Of2 res
      } |> Yield
    | Yield k1, Yield k2 -> (.||) k1 k2 |> Yield

let (.||>) s1 s2 = s1 .|| s2 |> ignore_
let concurrent = (.||)

let rec (==>) (c:Script<Option<'a>,'s>) (s:'a -> Script<'b,'s>) : Script<'b,'s> =
  script{
    let! x = c
    match x with
    | Some a ->
      do! arrow_yield_
      let! res = s a
      return res
    | None ->
      do! yield_
      return! (==>) c s
  }

let rec (=>) (c:Script<bool,'s>) (s:Script<'a,'s>) : Script<'a,'s> =
  script{
    let! x = c
    if x then
      do! arrow_yield_
      let! res = s
      return res
    else
      do! yield_
      return! (=>) c s
  }

let (=>>) s1 s2 = s1 => s2 |> ignore_
let guard_ = (=>)

let rec (!!) (s:Script<Unit,'s>) : Script<Unit,'s> =
  script{
    do! s
    return! (!!) s
  }
let repeat_ = (!!)

let rec (.|||) (ss:list<Script<'a,'b>>) : Script<'a,'b> =
  fun s ->
    let ss' = ss |> List.map(fun elem -> elem s)
    match ss' |> List.tryFind(function Return x -> true |_-> false) with
    |Microsoft.FSharp.Core.Some(Return x) -> Return x
    |Microsoft.FSharp.Core.None -> 
      match ss' |> List.tryFind(function ArrowYield k -> true |_-> false) with
      | Microsoft.FSharp.Core.Some(ArrowYield k) -> Yield k
      | Microsoft.FSharp.Core.None -> Yield((.|||) (ss' |> List.map (function Yield k -> k)))

let concurrent_many_first = (.|||)

let rec (!?) (p:Script<'a,'s>) : Script<'a,'s> =
  fun s ->
    match p s with
    | Return x -> Return x
    | Yield k | ArrowYield k -> (!?) k s

let atomic_ = (!?)

let time : Script<DateTime,'s> = fun _ -> Return(DateTime.Now)

let wait_doing (action:float32 -> Script<Unit,'s>) (interval:float32) : Script<Unit,'s> =
  script{
    let! t0 = time
    let rec wait() =
      script{
          let! t = time
          let dt = (t - t0).TotalSeconds |> float32
          if dt < interval then
            do! yield_
            do! action dt
            return! wait()
      }
    do! wait()
  }

let wait interval = wait_doing (fun (dt:float32) -> script{ return () }) interval

let game_pattern (init:Script<'a,'s>) 
                 (intro:Script<Unit,'s>)
                 (game_over:'a -> Script<Option<'b>,'s>)
                 (logic:'a -> Script<Unit,'s>) 
                 (ending:'a -> 'b -> Script<'c,'s>) : Script<'c,'s> =
  script{
    let! x = init |> (!?)
    do! intro
    let! (Choice1Of2 res) = ((game_over x) ==> ending x) .|| ((logic x) |> repeat_)
    return res
  }

let wait_condition p =
  let rec wait() =
    script{
      do! yield_
      if p() then
        return ()
      else
        return! wait()
    }
  wait()

let script_yield (): Script<Unit,'s> =
  script{
    do! yield_
  }
//let wait_game_over (game_over:Script<bool,'s>) : Script<Unit,'s> =
//  let null_script = script{ return () }
//  game_pattern null_script
//               (fun () -> game_over)
//               (fun () -> null_script)
//               (fun () -> null_script)
